from typing import Literal

from ..pipeline.pipeline import PipeModule, Port, status_manage
from ..pipeline.portTypes import PortTypeHint, PortType
from ..pipeline.pipeData import RangeModel, NestedRangeModel
from ..llm.dataExtractor import TableExtractionConfig, ExtractTableDataByImages
from ..dataclass.tables import TableData, TableCollection


@status_manage
class TableDataExtractor(PipeModule):
    """Extract table data from images using LLM"""

    def __init__(self, mname: str | None = "TableDataExtractor", auto_run: bool = True,
                 images_path: PortTypeHint.FilesPath | None = None,
                 images_prefix: str | list[str] | None = None,
                 columns: list[str] | dict[str, list[str]] | None = None,
                 table_description: str | dict[str, str] = "",
                 extraction_instructions: str | dict[str, str] | None = None,
                 data_types: dict[str, Literal["string", "number", "float", "int", "date", "boolean"]] | dict[str, dict[str, Literal["string", "number", "float", "int", "date", "boolean"]]] | None = None,
                 vision_model: str = "Qwen/Qwen2.5-VL-32B-Instruct",
                 temperature: float = 0.1,
                 max_tokens: int = 4096,
                 api_key: str = None):
        """
        Initialize the table data extractor

        Parameters
        ----------
        images_path: PortTypeHint.FilesPath | None, default: None
            Path to the images to extract table data from.
            If the images_path are assinged by input port, the self.images_path will be overwritten by the input data.

        images_prefix: str | list[str] | None, default: None
            Prefix of the images to extract table data from.
            If None, all the images will be used.

        columns: list[str] | dict[str, list[str]] | None, default: None
            List of column names to extract or a dictionary mapping table names to column names.
            If dict, the keys are the images_prefix, the values are the column names.

        table_description: str | dict[str, str] | None, default: "table"
            Description of what kind of table this is
            If dict, the keys are the images_prefix, the values are the table descriptions.

        extraction_instructions: str | dict[str, str] | None, default: None
            Custom instructions for data extraction
            If dict, the keys are the images_prefix, the values are the extraction instructions.

        data_types: dict[str, Literal["string", "number", "float", "int", "date", "boolean"]] | dict[str, dict[str, Literal["string", "number", "float", "int", "date", "boolean"]]] | None, default: None
            Dictionary mapping column names to expected data types
            If dict, the keys are the images_prefix, the values are the data types.

        vision_model: str, default: "Qwen/Qwen2.5-VL-32B-Instruct"
            Vision model to use for extraction

        temperature: float, default: 0.1
            Model temperature for generation

        max_tokens: int, default: 4096
            Maximum tokens for model response

        api_key: str, default: None
            The API key for the LLM API.
            Only SiliconFlow API key is supported at current version.
        """
        super().__init__(mname=mname, auto_run=auto_run)

        self._ports_in = {"InputImages": Port(ptype=PortType.FilesPath, data=images_path, pdoc="The path of the images.")}
        self._ports_out = {"OutputTables": Port(ptype=[PortType.TableData, PortType.TableCollection], pdoc="The data of the tables.")}

        self.images_path = images_path
        self.images_prefix = images_prefix
        self.columns = columns
        self.table_description = table_description
        self.extraction_instructions = extraction_instructions
        self.data_types = data_types or {}
        self.vision_model = vision_model
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.api_key = api_key

    @property
    def InputImages(self) -> PortTypeHint.FilesPath | None:
        raise AttributeError("Property 'InputImages' is write-only.")
    
    @InputImages.setter
    def InputImages(self, value: PortTypeHint.FilesPath | None) -> None:
        self["InputImages"] = value
    
    @property
    def OutputTables(self) -> PortTypeHint.TableData | PortTypeHint.TableCollection | None:
        return self["OutputTables"]
    
    def set_cal_params(self, reset: bool = True) -> dict[str, RangeModel | NestedRangeModel]:
        if self.images_prefix is None:
            columns_range_model = RangeModel(vtype="list", list_type="str", title="列名", default=self.columns, depends_on="images_prefix")
            table_description_range_model = RangeModel(vtype="str", title="表描述", default=self.table_description, depends_on="images_prefix")
            extraction_instructions_range_model = RangeModel(vtype="str", title="补充提示词", default=self.extraction_instructions, depends_on="images_prefix")
            data_types_range_model = RangeModel(vtype="dict", dict_value_type="str", 
                                                title="各列数据类型", default=self.data_types, dict_key_choices=self.columns,
                                                dict_value_choices=["string", "number", "float", "int", "date", "boolean"], depends_on="images_prefix")
        else:
            if isinstance(self.images_prefix, str):
                images_prefix = [self.images_prefix]
            else:
                images_prefix = self.images_prefix
            columns_range_model = RangeModel(vtype="dict", dict_key_choices=images_prefix, dict_value_type="list", list_type="str", 
                                             title="列名", default=self.columns, depends_on="images_prefix")
            table_description_range_model = RangeModel(vtype="dict", dict_key_choices=images_prefix, dict_value_type="str", 
                                                       title="表描述", default=self.table_description, depends_on="images_prefix")
            extraction_instructions_range_model = RangeModel(vtype="dict", dict_key_choices=images_prefix, dict_value_type="str", 
                                                             title="补充提示词", default=self.table_description, depends_on="images_prefix")
            
            fields = {}
            for image_prefix in images_prefix:
                if self.columns is None:
                    choices = None
                elif isinstance(self.columns, list):
                    choices = self.columns
                else:
                    choices = self.columns[image_prefix]
                fields[image_prefix] = RangeModel(vtype="dict", dict_value_type="str",
                                                title="各列数据类型", default=self.data_types, 
                                                dict_key_choices=choices, dict_value_choices=["string", "number", "float", "int", "date", "boolean"])
            data_types_range_model = NestedRangeModel(vtype="dict", title="各列数据类型", fields=fields, default=self.data_types, 
                                                      depends_on="images_prefix", min_max_items=(len(images_prefix), len(images_prefix)))
            

        return {"images_path": RangeModel(vtype=("str", "list"), list_type="str", title="图片路径", default=self.images_path),
                "images_prefix": RangeModel(vtype=("str", "list", "None"), list_type="str", title="图片前缀", default=self.images_prefix),
                "columns": columns_range_model,
                "table_description": table_description_range_model,
                "extraction_instructions": extraction_instructions_range_model,
                "data_types": data_types_range_model,
                "vision_model": RangeModel(vtype="str", title="视觉模型", default=self.vision_model)}
    
    def execute(self) -> PortTypeHint.GeneralTable | None:
        images_path = self._ports_in["InputImages"].data
        if (images_path is None and self.images_path is None) or self.columns is None:
            self._ports_out["OutputTables"].data = None
            return None
        
        if images_path is None:
            images_path = self.images_path
        else:
            self.images_path = images_path
        
        # If images_prefix is None or string, extract a single table
        if self.images_prefix is None or isinstance(self.images_prefix, str):
            # Get configuration values for single table extraction
            if isinstance(self.images_prefix, str):
                prefix = self.images_prefix
                # Get configuration values for this prefix
                if isinstance(self.columns, dict):
                    columns = self.columns.get(prefix, self.columns)
                else:
                    columns = self.columns
                
                if isinstance(self.table_description, dict):
                    table_description = self.table_description.get(prefix, "")
                else:
                    table_description = self.table_description
                
                if isinstance(self.extraction_instructions, dict):
                    extraction_instructions = self.extraction_instructions.get(prefix, None)
                else:
                    extraction_instructions = self.extraction_instructions
                
                if isinstance(self.data_types, dict) and prefix in self.data_types:
                    data_types = self.data_types[prefix]
                else:
                    data_types = self.data_types
            else:
                # Use default values when images_prefix is None
                columns = self.columns
                table_description = self.table_description
                extraction_instructions = self.extraction_instructions
                data_types = self.data_types
                prefix = None
            
            extraction_config = TableExtractionConfig(columns=columns, table_description=table_description, 
                                                      extraction_instructions=extraction_instructions, 
                                                      data_types=data_types, vision_model=self.vision_model, 
                                                      temperature=self.temperature, max_tokens=self.max_tokens)
            
            table_data_extractor = ExtractTableDataByImages(images_path=images_path, extraction_config=extraction_config, 
                                                           api_key=self.api_key, images_prefix=prefix)
            table_data = table_data_extractor.process()
            table_data = TableData(data=table_data, name=f"Table_{prefix}", title=f"Table_{prefix}", description=table_description)

            self._ports_out["OutputTables"].data = table_data
            return table_data
        
        # Handle multiple table extraction with images_prefix (list)
        prefixes = self.images_prefix
        
        # Create TableCollection to store multiple tables
        table_collection = TableCollection(name="ExtractedTables", title="ExtractedTables", description="Extracted Tables from Images")
        
        # Extract table for each prefix
        for prefix in prefixes:
            # Get configuration values for this prefix
            if isinstance(self.columns, dict):
                columns = self.columns.get(prefix, self.columns)
            else:
                columns = self.columns
            
            if isinstance(self.table_description, dict):
                table_description = self.table_description.get(prefix, "")
            else:
                table_description = self.table_description
            
            if isinstance(self.extraction_instructions, dict):
                extraction_instructions = self.extraction_instructions.get(prefix, None)
            else:
                extraction_instructions = self.extraction_instructions
            
            if isinstance(self.data_types, dict) and prefix in self.data_types:
                data_types = self.data_types[prefix]
            else:
                data_types = self.data_types
            
            # Create extraction config for this prefix
            extraction_config = TableExtractionConfig(columns=columns, table_description=table_description, 
                                                      extraction_instructions=extraction_instructions, 
                                                      data_types=data_types, vision_model=self.vision_model, 
                                                      temperature=self.temperature, max_tokens=self.max_tokens)
            
            # Extract table data for this prefix
            table_data_extractor = ExtractTableDataByImages(images_path=images_path, extraction_config=extraction_config, 
                                                           api_key=self.api_key, images_prefix=prefix)
            table_data = table_data_extractor.process()
            table_data = TableData(data=table_data, name=f"Table_{prefix}", title=f"Table_{prefix}", description=table_description)
            
            # Add extracted table to collection
            table_collection.add_table(table_data)

        self._ports_out["OutputTables"].data = table_collection
        return table_collection

